from torch import nn

class policy_selector(nn.Module):
    def __init__(self,input_size,
                 hidden_size):
        super(policy_selector, self).__init__()
        self.affine1 = nn.Linear(input_size, hidden_size)
        self.affine2 = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()
        self.saved_log_probs = []
        self.rewards = []

    def forward(self, x):
        x = self.affine1(x)
        action_scores = self.affine2(x)
        return self.sigmoid(action_scores)